package com.jzlin.funhub.admin.service.chatai.store;

import cn.hutool.core.collection.CollUtil;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 内存存储.
 *
 * @author linjz
 * @since 2024/11/12
 */
public class InMemoryChatStore implements IChatStore {
    Map<String, List<ChatAiMessageRecord>> context = new ConcurrentHashMap<>();

    /**
     * 获取会话上下文消息
     *
     * @param sessionId 会话id
     * @param count     关联数量
     * @return 上下文消息列表
     */
    @Override
    public List<ChatAiMessageRecord> getRefContextMessage(String sessionId, int count) {
        List<ChatAiMessageRecord> records = context.get(sessionId);
        if (CollUtil.isEmpty(records)) {
            return Collections.emptyList();
        }
        if (count >= records.size()) {
            return records;
        }

        return CollUtil.sub(records, records.size() - count, records.size());
    }

    /**
     * 保存消息
     *
     * @param record 消息
     * @return 成功返回true
     */
    @Override
    public boolean saveMessage(ChatAiMessageRecord record) {
        return context.computeIfAbsent(record.getSessionId(), (k) -> new ArrayList<>()).add(record);
    }
}
